package zuo.algo.c12;

import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

/**
 * 00:25-并查集
 */
public class UnionFind {

    public static class Element<V> {

        public V value;

        public Element(V value) {
            this.value = value;
        }
    }

    public static class UnionFindSet<V> {

        public Map<V, Element<V>> elementMap;
        public Map<Element<V>, Element<V>> parentMap;
        public Map<Element<V>, Integer> sizeMap;

        public UnionFindSet(List<V> valueList) {
            elementMap = new HashMap<>(valueList.size());
            parentMap = new HashMap<>(valueList.size());
            sizeMap = new HashMap<>(valueList.size());

            for (V value : valueList) {
                Element<V> element = new Element<>(value);
                elementMap.put(value, element);
                parentMap.put(element, element);
                sizeMap.put(element, 1);
            }
        }

        public boolean isSameSet(V a, V b) {
            Element<V> element1 = elementMap.get(a);
            Element<V> element2 = elementMap.get(b);
            if (element1 == null || element2 == null) {
                return false;
            }
            Element<V> head1 = findHead(element1);
            Element<V> head2 = findHead(element2);
            return head1 == head2;
        }

        public void union(V a, V b) {
            Element<V> element1 = elementMap.get(a);
            Element<V> element2 = elementMap.get(b);
            if (element1 == null || element2 == null) {
                return;
            }

            Element<V> head1 = findHead(element1);
            Element<V> head2 = findHead(element2);
            if (head1 == head2) {
                return;
            }

            Integer size1 = sizeMap.get(head1);
            Integer size2 = sizeMap.get(head2);
            if (size1 > size2) {
                parentMap.put(head2, head1);
                sizeMap.put(head1, size1 + size2);
                sizeMap.remove(head2);
            } else {
                parentMap.put(head1, head2);
                sizeMap.put(head2, size1 + size2);
                sizeMap.remove(head1);
            }
        }

        private Element<V> findHead(Element<V> element) {
            Element<V> pElement = parentMap.get(element);
            LinkedList<Element<V>> queue = new LinkedList<>();
            while (pElement != element) {
                queue.add(element);
                element = pElement;
                pElement = parentMap.get(element);
            }

            while (!queue.isEmpty()) {
                parentMap.put(queue.poll(), element);
            }
            return element;
        }
    }
}
